import subprocess
import re
import os


# 这是真的
class func_node:
    def __init__(self):
        self.func_name = ""
        self.func_lib = ""
        self.func_bind = ""
        self.visibility = ""
        self.depend_funcs = []


class SoAnalyser:
    @staticmethod
    def get_so_contents(file_path):
        if ".txt" in file_path:
            with open(file_path, 'r', encoding="utf-8") as file:
                # 读取文件内容
                return file.read()
        cmd = f"objdump -d {file_path}"
        try:
            output = subprocess.check_output(cmd, shell=True).decode()
        except Exception:
            return ""
        return output

    # 分析so文件的符号表信息，从中分离出此文件provide和require的函数（只有内部使用的就忽略了）
    @staticmethod
    def analyse_so_outer_funcs(file_path):
        func_dict = {"require":[], "provide":[]}
        try:
            cmd = f"readelf -s {file_path}"
            sym_table = subprocess.check_output(cmd, shell=True).decode()
        except Exception:
            return func_dict
        sym_table = sym_table.split("\n")[3:]
        for sym in sym_table:
            sym_val = sym.split(" ")
            sym_val[:] = [item for item in sym_val if item != '']
            try:
                sym_num = sym_val[0]
                sym_value = sym_val[1]
                sym_size = sym_val[2]
                sym_type = sym_val[3]
                sym_bind = sym_val[4]
                sym_vis = sym_val[5]
                sym_ndx = sym_val[6]
                sym_name = sym_val[7]
            except Exception:
                continue

            if sym_type == "FUNC" and (sym_bind == "GLOBAL" or sym_bind == "WEAK"):
                if sym_ndx == "UND":
                    func_dict["require"].append(sym_name.split("@")[0])
                else :
                    func_dict["provide"].append(sym_name.split("@")[0])
        return func_dict


    # 除了函数名@@库名.版本号和函数名@库名的格式外，还有一些其他的函数名格式，例如：
    #
    #     函数名：没有库名和版本号信息，表示该函数在程序的可执行文件中。
    #     函数名@版本号：只有版本号信息，没有库名信息。
    #     函数名@@库名：没有版本号信息，只有库名信息。
    #
    # 需要根据具体的反汇编工具和二进制文件格式来确定函数名的具体格式
    @staticmethod
    def analyse_so(file_path):
        # 获取当前so文件要从外部获取的函数名信息
        # so_require_funcs = SoAnalyser.analyse_so_outer_funcs(file_path)["require"]
        l = SoAnalyser.get_so_contents(file_path).split("\n")

        func_name = ""
        func_id = ""
        call_dict = {}
        # 定义正则表达式
        pattern = r'^([0-9a-fA-F]+)\s+<(\S+)@+([^\s>]+)>:$'
        pattern2 = r'\s*\w+:\s*\w+\s+\w+\s+(.*callq.*)\s+<([^>]+)>'

        for i in range(len(l)):
            sentence = l[i]
            match = re.match(pattern, sentence)
            # 输出匹配结果
            if match:
                # 这个函数名应该就是转跳转到plt指定位置，是一段辅助链接的语句，不是在文件中定义的函数
                # 函数应该是来源于外部的
                # 这个判断条件足够严谨吗？还要斟酌一下
                # if "plt" in sentence and not (i+2 < len(l) and "jmpq" in l[i+1] and "xchg" in l[i+2]):
                #     print(sentence)
                # if i+2 < len(l) and "jmpq" in l[i+1] and "xchg" in l[i+2]:
                #     #print(sentence)
                #     #print(l[i+1])
                #     #print(l[i+2])
                #     if "plt" not in sentence:
                #         print("+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++")
                #     continue
                # 上面主要是一段用于检验的代码，得到的结论是：
                #       有些看似是函数名，但实际作用就是动态链接的跳转，是一段辅助链接的语句，不是在文件中定义的函数
                #       这些函数应该是来源于外部的，相当于外部函数的声明
                #       这个函数一定只包含在plt中的
                #       有plt名的函数定义的一定是这类型的“伪函数”
                if i + 2 < len(l) and "jmpq" in l[i + 1] and "xchg" in l[i + 2]:
                    # 这是外部函数的声明
                    continue
                    # print(sentence)
                    # print(l[i+1])
                    # print(l[i+2])
                func_name = match.group(2).replace("@", "")
                belong_lib = match.group(3).replace("@", "")
                func_id = func_name + "@" + belong_lib
                call_dict[func_id] = {}
                #print('地址:', match.group(1))
                # print('函数名:', func_name)
                # print('函数所属库和版本号:', belong_lib)
                continue
            match2 = re.match(pattern2, sentence)
            if match2:
                if func_id == "":
                    continue
                depend_id = match2.group(2)
                depend_lib = depend_id.split("@")[-1]
                depend_func = (depend_id.split("@")[0])
                if depend_id not in call_dict[func_id]:
                    call_dict[func_id][depend_id] = 0
                call_dict[func_id][depend_id] = call_dict[func_id][depend_id] + 1
                # call_func_list.append(match2.group(2))
                # print(match2.group(0))
                # print('      调用函数名:', match2.group(2))
        return call_dict
        # with open(output_path, "w", encoding='utf-8') as f1:  ## 设置'utf-8'编码
        #     f1.write(json.dumps(call_dict,
        #                         ensure_ascii=False,
        #                         indent=4,
        #                         separators=(',', ': ')))  # (先转换成字符串再存储)

    # 给出两个文件包（经过处理的，一级目录下就包括了所有so文件的），获取依赖的所有函数
    @staticmethod
    def find_depend_func_num(provide_pkg_dir, require_pkg_dir):
        provide_set = set()
        provide_so_files = os.listdir(provide_pkg_dir)
        for file in provide_so_files:
            file_path = os.path.join(provide_pkg_dir, file)
            provide_pkg_funcs = SoAnalyser.analyse_so_outer_funcs(file_path)["provide"]
            provide_set = provide_set.union(set(provide_pkg_funcs))

        require_set = set()
        require_so_files = os.listdir(require_pkg_dir)
        for file in require_so_files:
            file_path = os.path.join(require_pkg_dir, file)
            require_pkg_funcs = SoAnalyser.analyse_so_outer_funcs(file_path)["require"]
            require_set = require_set.union(set(require_pkg_funcs))

        depend_funcs_set = provide_set.intersection(require_set)
        return depend_funcs_set

    # 给出两个文件包（经过处理的，一级目录下就包括了所有so文件的），获取函数调用关系图
    @staticmethod
    def find_depend_graph(provide_pkg_id, require_pkg_id):
        provide_pkg_dir = "./" + provide_pkg_id
        require_pkg_dir = "./" + require_pkg_id
        depend_graph = {}

        provide_set = set()
        provide_call_graph = {}
        provide_so_files = os.listdir(provide_pkg_dir)
        for file in provide_so_files:
            file_path = os.path.join(provide_pkg_dir, file)
            call_graph = SoAnalyser.analyse_so(file_path)
            # 会不会有重复？应该不会？但需要验证一下
            provide_call_graph.update(call_graph)
            provide_pkg_funcs = SoAnalyser.analyse_so_outer_funcs(file_path)["provide"]
            provide_set = provide_set.union(set(provide_pkg_funcs))
        # print(len(provide_call_graph))
        provide_name2id_dict = {}
        for provide_func_id in provide_call_graph:
            provide_func_name = provide_func_id.split("@")[0]
            if provide_func_name in provide_set:
                # 这里可能存在重复，导致覆盖，不准确，还需验证
                provide_name2id_dict[provide_func_name] = provide_func_id

        require_call_graph = {}
        require_set = set()
        require_so_files = os.listdir(require_pkg_dir)
        print(f'{require_pkg_id}: ')
        for file in require_so_files:
            file_path = os.path.join(require_pkg_dir, file)
            call_graph = SoAnalyser.analyse_so(file_path)
            # 会有重复，比如对于同一个包中不同so文件的内部函数
            require_call_graph.update(call_graph)
            print(f"{require_pkg_id}#{file}: ")
            for func_name in call_graph.keys():
                print(f"  func id: {func_name}: ")
                depend_graph[f"{require_pkg_id}#{file}#{func_name}"] = []
                depend_outer_list = []
                require_funcs_set = set(call_graph[func_name].keys())
                for require_func_id in require_funcs_set:
                    require_func_name = require_func_id.split("@")[0]
                    if require_func_name in provide_name2id_dict:
                        print(f"      {provide_pkg_id}#{provide_name2id_dict[require_func_name]}")
                        depend_outer_list.append(f'{provide_pkg_id}#{provide_name2id_dict[require_func_name]}')
                if depend_outer_list == []:
                    print("      None")
                depend_graph[f"{require_pkg_id}#{file}#{func_name}"] = depend_outer_list
        return depend_graph